import os
from tqdm import tqdm
from PIL import Image, ImageDraw
import numpy as np

current_path = os.path.split(os.path.abspath(__file__))[0]
# Load the images


val_index = 1471

with open('train.txt', 'w') as train_txt:
    with open('val.txt', 'w') as val_txt:
        for image_path in tqdm(sorted(os.listdir('mask'))):

            # 如果大于某个序号，就都保存在验证集的txt中
            if int(os.path.splitext(image_path)[0]) < val_index:
                train_txt.write(os.path.join(current_path, 'images', image_path) + '\n')
            else:
                val_txt.write(os.path.join(current_path, 'images', image_path) + '\n')
            mask = Image.open(os.path.join('mask', image_path))

            # 获取颜色
            colors = Image.Image.getcolors(mask)
            colors = [i[1] for i in colors if i[1] != (0, 0, 0)]
            print(colors)

            # Convert images to numpy array
            mask_np = np.array(mask)

            with open('labels/' + os.path.splitext(image_path)[0] + '.txt', 'w') as f:

                for color in colors:

                    # 获取不同颜色区域的像素坐标
                    non_black_pixels = np.where(np.all(mask_np == color, axis=-1))

                    # 通过像素坐标，得到x,y,w,h的标注
                    x, y = np.min(non_black_pixels[1]), np.min(non_black_pixels[0])
                    w, h = np.max(non_black_pixels[1]) - x, np.max(non_black_pixels[0]) - y

                    print(f"Bounding box coordinates: x={x}, y={y}, width={w}, height={h}")

                    # yolo的标注格式的x,y是中心点，所以要做一下转换
                    yolo_x = (x + w / 2) / mask_np.shape[0]
                    yolo_y = (y + h / 2) / mask_np.shape[1]
                    yolo_w = w / mask_np.shape[0]
                    yolo_h = h / mask_np.shape[1]
                    f.write(f'0 {yolo_x} {yolo_y} {yolo_w} {yolo_h}\n')
                f.close()


            ## 验证图片
            #     draw = ImageDraw.Draw(mask)
            #     #绘制矩形
            #     draw.rectangle((x,y,x+w,y+h),outline=(255,0,0))
            #     #绘制正方形
            #
            # mask.save("draw_img.jpg")

# 最后自己生成 yaml，用于告知地址

